import matplotlib.pyplot as plt
import numpy as np


classes = ("airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck")

"""
绘制平均准确率
"""
plt.figure()
with open('../results/total.txt', 'r') as f:
    data = f.readlines()
    data = [float(i.strip('\n')) for i in data]
    data_num = [i for i in range(len(data))]
my_x_ticks = np.arange(0, 512, 50)
my_y_ticks = np.arange(0, 100, 5)
plt.plot(data_num, data)
plt.grid()
plt.xticks(my_x_ticks)
plt.yticks(my_y_ticks)
plt.xlabel("Prune K Neuron(s)")
plt.ylabel("Total Acc")
plt.title("Changing Total Accuracy with K")
plt.show()

"""
绘制各类的准确率
"""
plt.figure()
for idx in range(1, 1 + len(classes)):
    plt.subplot(4, 3, idx)
    with open('../results/{}.txt'.format(classes[idx - 1]), 'r') as f:
        data = f.readlines()
        data = [float(i.strip('\n')) for i in data]
        data_num = [i for i in range(len(data))]
        # my_x_ticks = np.arange(0, 512, 50)
        # my_y_ticks = np.arange(0, 100, 5)
        plt.plot(data_num, data, label=classes[idx - 1])
        plt.legend(loc='lower left', prop={'size': 7})
        # plt.grid()
        plt.xticks([])
        plt.yticks([])
        # plt.xticks(my_x_ticks)
        # plt.yticks(my_y_ticks)
        # plt.xlabel("Prune K Neuron(s)")
        # plt.ylabel("Acc")
        # plt.title("Changing {} Accuracy with K".format(classes[idx - 1]))
plt.suptitle('Changing Accuracy with K')
plt.show()
